from setting import backtest
from setting import data
import polars as pl
import datetime
import pickle
from setting.crate import *

# 策略名称 - 用于生成测试报告
strategy_name = "lyy"

# 合约配置
base_data_symbol = ["SHFE_ag00_300"]  # 基准合约列表(用于时间轴)
trade_data_symbol = [
    "CZCE_AP00_300", "CZCE_CF00_300", "CZCE_CJ00_300", "CZCE_CY00_300", "CZCE_FG00_300",
    "CZCE_LR00_300", "CZCE_MA00_300", "CZCE_OI00_300", "CZCE_PF00_300", "CZCE_PK00_300",
    "CZCE_PL00_300", "CZCE_PM00_300", "CZCE_PR00_300", "CZCE_RI00_300", "CZCE_RM00_300",
    "CZCE_RS00_300", "CZCE_SA00_300", "CZCE_SH00_300", "CZCE_SF00_300", "CZCE_SM00_300",
    "CZCE_SR00_300", "CZCE_TA00_300", "CZCE_UR00_300", "CZCE_WH00_300", "CZCE_ZC00_300",
    "DCE_a00_300", "DCE_b00_300", "DCE_bb00_300", "DCE_bz00_300", "DCE_c00_300",
    "DCE_cs00_300", "DCE_eb00_300", "DCE_eg00_300", "DCE_fb00_300", "DCE_j00_300",
    "DCE_jd00_300", "DCE_jm00_300", "DCE_l00_300", "DCE_lg00_300", "DCE_lh00_300",
    "DCE_m00_300", "DCE_p00_300", "DCE_pg00_300", "DCE_pp00_300", "DCE_rr00_300",
    "DCE_v00_300", "DCE_y00_300", "GFEX_lc00_300", "GFEX_ps00_300", "GFEX_si00_300",
    "INE_bc00_300", "INE_ec00_300", "INE_lu00_300", "INE_nr00_300", "INE_sc00_300",
    "SHFE_ad00_300", "SHFE_ag00_300", "SHFE_al00_300", "SHFE_ao00_300", "SHFE_br00_300",
    "SHFE_bu00_300", "SHFE_cu00_300", "SHFE_fu00_300", "SHFE_hc00_300", "SHFE_ni00_300",
    "SHFE_op00_300", "SHFE_pb00_300", "SHFE_rb00_300", "SHFE_ru00_300", "SHFE_sn00_300",
    "SHFE_sp00_300", "SHFE_ss00_300", "SHFE_wr00_300", "SHFE_zn00_300"
]

# 跨周期数据合约池(日线数据)
extra_data_symbol = [
    "CZCE_AP00_86400", "CZCE_CF00_86400", "CZCE_CJ00_86400", "CZCE_CY00_86400", "CZCE_FG00_86400",
    "CZCE_LR00_86400", "CZCE_MA00_86400", "CZCE_OI00_86400", "CZCE_PF00_86400", "CZCE_PK00_86400",
    "CZCE_PL00_86400", "CZCE_PM00_86400", "CZCE_PR00_86400", "CZCE_RI00_86400", "CZCE_RM00_86400",
    "CZCE_RS00_86400", "CZCE_SA00_86400", "CZCE_SH00_86400", "CZCE_SF00_86400", "CZCE_SM00_86400",
    "CZCE_SR00_86400", "CZCE_TA00_86400", "CZCE_UR00_86400", "CZCE_WH00_86400", "CZCE_ZC00_86400",
    "DCE_a00_86400", "DCE_b00_86400", "DCE_bb00_86400", "DCE_bz00_86400", "DCE_c00_86400",
    "DCE_cs00_86400", "DCE_eb00_86400", "DCE_eg00_86400", "DCE_fb00_86400", "DCE_j00_86400",
    "DCE_jd00_86400", "DCE_jm00_86400", "DCE_l00_86400", "DCE_lg00_86400", "DCE_lh00_86400",
    "DCE_m00_86400", "DCE_p00_86400", "DCE_pg00_86400", "DCE_pp00_86400", "DCE_rr00_86400",
    "DCE_v00_86400", "DCE_y00_86400", "GFEX_lc00_86400", "GFEX_ps00_86400", "GFEX_si00_86400",
    "INE_bc00_86400", "INE_ec00_86400", "INE_lu00_86400", "INE_nr00_86400", "INE_sc00_86400",
    "SHFE_ad00_86400", "SHFE_ag00_86400", "SHFE_al00_86400", "SHFE_ao00_86400", "SHFE_br00_86400",
    "SHFE_bu00_86400", "SHFE_cu00_86400", "SHFE_fu00_86400", "SHFE_hc00_86400", "SHFE_ni00_86400",
    "SHFE_op00_86400", "SHFE_pb00_86400", "SHFE_rb00_86400", "SHFE_ru00_86400", "SHFE_sn00_86400",
    "SHFE_sp00_86400", "SHFE_ss00_86400", "SHFE_wr00_86400", "SHFE_zn00_86400"
]

# 加载数据
print('加载前', datetime.datetime.now().strftime("%H:%M:%S"))
base_data, all_trade_data = data.get_data(
    base_data_symbol, 
    trade_data_symbol + extra_data_symbol, 
    "20250101",
    "20251001"
)

# 分割数据
trade_data = {k: v for k, v in all_trade_data.items() if k in trade_data_symbol}
extra_data = {k: v for k, v in all_trade_data.items() if k in extra_data_symbol}
print('加载后', datetime.datetime.now().strftime("%H:%M:%S"))

# 回测参数配置
initial_cash = 10000 * 10000.0  # 初始资金:1000万
fee_rate = 3.0 / 10000  # 手续费率:万分之三

# 加载合约参数
with open("setting/volume_multiple.pkl", "rb") as f:
    loaded_volume = pickle.load(f)
with open("setting/margin_ratio.pkl", "rb") as f:
    loaded_margin = pickle.load(f)

deposit_rates = loaded_margin
multipliers = loaded_volume

# 处理额外数据 - 添加涨幅和相关指标
for symbol, data in extra_data.items():
    extra_data[symbol] = data.lazy().with_columns([
        # 基础涨幅
        ((pl.col("close") - pl.col("close").shift(1)) / pl.col("close").shift(1)).alias("pct_change"),
        # 振幅:(最高价 - 最低价) / 前收盘
        ((pl.col("high") - pl.col("low")) / pl.col("close").shift(1)).alias("amplitude"),
    ]).collect()
class Strategy(PythonStrategy):
    """
    动量策略:持有涨幅前5的品种
    """

    def __init__(self, backtest_instance, trade_data, extra_data):
        super().__init__(backtest_instance, trade_data, extra_data)

    def on_bar(self, timestamp):
        """K线回调函数 - 每个时间戳调用一次"""
        symbol_list = {}
        data = {}
        
        for symbol in self.trade_data.keys():
            recent_data = self.get_recent_data(symbol, timestamp, lookback=2, require_exact_match=True)
            recent_data_daily = self.get_recent_data(extract_symbol_base(symbol) + "_86400", timestamp, lookback=2, require_exact_match=False)
            if recent_data is not None and recent_data_daily is not None:
                pct_change = recent_data_daily["pct_change"][-1]
                if pct_change is not None:
                    symbol_list[extract_symbol_base(symbol)] = recent_data_daily["pct_change"][-1]
                    data[extract_symbol_base(symbol)] = recent_data
                
        self.calculate_momentum_signal(symbol_list, timestamp, data)

    def calculate_momentum_signal(self, symbol_pct_changes, timestamp, data):
        """动量信号计算函数 - 持有涨幅前5的品种"""
        if not symbol_pct_changes or len(symbol_pct_changes) < 5:  # 至少要有5个品种
            return
        # 按涨幅降序排序(涨幅大的在前)
        sorted_symbols = sorted(symbol_pct_changes.items(), key=lambda x: x[1], reverse=True)
        
        # 取前5名(涨幅最大的5个)
        top_5_symbols = [symbol for symbol, _ in sorted_symbols[:5]]
        
        # 获取当前持仓
        current_long_positions = set(self.backtest_instance.long_position.keys())
        
        # 计算还可以买入多少个品种
        MAX_POSITIONS = 5  # 最大持仓品种数量
        current_position_count = len(current_long_positions)
        available_slots = MAX_POSITIONS - current_position_count
        
        # 买入逻辑:买入前5名的品种(如果还没有持仓且还有仓位空间)
        for symbol in top_5_symbols:
            if available_slots <= 0:
                break  # 没有仓位空间了，停止买入
            if symbol not in current_long_positions:
                current_price = data[symbol]["close"][-1]
                # 买入1手
                self.backtest_instance.buy(symbol, 1, current_price)
                available_slots -= 1  # 减少可用仓位
        
        # 卖出逻辑:卖出不在前5名的持仓
        for symbol in list(current_long_positions):
            if symbol not in data.keys():
                continue
            if symbol not in top_5_symbols:
                recent_data = data[symbol]
                if recent_data is not None and len(recent_data) > 0:
                    current_price = recent_data["close"][-1]
                    # 获取当前持仓数量
                    current_qty = self.backtest_instance.long_position.get(symbol, 0)
                    if current_qty > 0:
                        self.backtest_instance.sell(symbol, current_qty, current_price)


def main():
    """
    主函数 - 配置并运行回测
    """
    # 初始化回测引擎
    my_backtest = backtest.BackTest(initial_cash, fee_rate, deposit_rates, multipliers, trade_data)

    # 创建策略实例 - 使用子类Strategy
    strategy = Strategy(my_backtest, trade_data, extra_data)

    # 准备回测时间序列
    datetime_strings = base_data.get_column("datetime").dt.strftime("%Y-%m-%d %H:%M:%S")

    # 启动回测
    backtest.start_backtest(
        strategy,
        my_backtest,
        datetime_strings,
        base_data.get_column("timestamps"),
        strategy_name
    )


if __name__ == "__main__":
    main()